#参考文章https://blog.csdn.net/hxxjxw/article/details/124071198
import torch
from torch import nn

input = torch.randn(8, 3, 20)
linear = nn.Linear(20, 40, bias=False)
wn_layer = nn.utils.weight_norm(linear, name='weight')
wn_output = wn_layer(input)

weight_direction = linear.weight / torch.norm(linear.weight, p=2, dim=1, keepdim=True)  # 二范数
weight_magnitude = wn_layer.weight_g
output = input @ (weight_direction.permute(1, 0).contiguous() * weight_magnitude.permute(1, 0).contiguous())
assert torch.allclose(wn_output, output)